"""
内容：图像预处理
日期：2020年7月10日
作者：Howie
"""
import numpy as np
from keras.utils import to_categorical
from keras.preprocessing.image import ImageDataGenerator


def load_mnist():
    """
    # 生成数据集
    :return: 划分好的训练集和测试集
    """
    with np.load('../dataset/mnist/mnist.npz') as data:
        (X_train, Y_train), (X_test, Y_test) = (
            data['x_train'], data['y_train']), (data['x_test'], data['y_test'])
    X_train = X_train.reshape(60000, 28, 28, 1).astype('float32') / 255.0
    X_test = X_test.reshape(10000, 28, 28, 1).astype('float32') / 255.0
    Y_train = to_categorical(Y_train, 10)
    Y_test = to_categorical(Y_test, 10)
    # 通过实时数据增强生成张量图像数据批次
    train_datagen = ImageDataGenerator(rescale=1. / 255,
                                       rotation_range=10,
                                       width_shift_range=0.2,
                                       height_shift_range=0.2,
                                       shear_range=0.7,
                                       zoom_range=[0.9, 2.2],
                                       horizontal_flip=True,
                                       vertical_flip=True,
                                       fill_mode='nearest')
    test_datagen = ImageDataGenerator(rescale=1. / 255)
    train_generator = train_datagen.flow(X_train, Y_train, batch_size=64)
    test_generator = test_datagen.flow(X_test, Y_test, batch_size=64)

    print(
        "train on {} samples, test on {} samples.".format(
            train_generator.x.shape[0],
            test_generator.x.shape[0]))



if __name__ == '__main__':
    load_mnist()
